/*
 * Copyright 2022 NXP
 * All rights reserved.
 *
 * SPDX-License-Identifier: BSD-3-Clause
 */

#include <math.h>
#include <stdlib.h>
#include <string.h>

#include "mel.h"

#define PI 3.1415926f

float MelScale(float freq) {
    return 1127.0f * logf(1.0f + freq / 700.0f);
}

float InverseMelScale(float mel) {
    return 700.0f * (expf(mel / 1127.0f) - 1.0f);
}

float window_func[FRAME_LEN];
void init_window_func(){
    for(int i=0;i<FRAME_LEN;i++){
        window_func[i] = 0.5f - 0.5f * cosf(2 * PI * i / FRAME_LEN);
    }
}

#include "arm_math.h"
static arm_rfft_fast_instance_f32 rfft;
void init_MelSG(){
    // init the fft    
    arm_rfft_fast_init_f32(&rfft, FRAME_LEN);
    
    init_window_func();
}

typedef struct {
    float real;
    float imag;
}complex_t;

static inline uint8_t quantize_melSG(float mel_f){
    float mel_clip;
    arm_clip_f32(&mel_f, &mel_clip, MIN_DB, MAX_DB, 1);
    mel_clip -= MIN_DB;
    mel_clip *= 255.5f / (MAX_DB - MIN_DB);
    return (uint8_t)mel_clip;   
}
void MelSG_compute_fast(int16_t* audio_data, uint32_t audio_len, uint8_t *melSG_out) {
    // 1. TensorFlow way of normalizing .wav data to (-1, 1), and apply window function
    float f_audio[FRAME_LEN];
    
    complex_t rfftAry[(FRAME_LEN>>1) + 1]; // 129, half fft, real-complex
    float buffer[(FRAME_LEN>>1) + 1]; // 129, including the 0, current
    int i = 0;
    for (; i < audio_len; i++) {
        f_audio[i] = audio_data[i] / 32768.0f * window_func[i];
    }
    for (; i < FRAME_LEN; i++) {
        f_audio[i] = 0;
    }

    // 3. Compute FFT
    arm_rfft_fast_f32(&rfft, f_audio, (float*)rfftAry, 0);

    // 4. Convert to power spectrum square root
    // uses self.buffer to store power of each fft bin
    for (int i = 1; i < ((FRAME_LEN>>1)) - 1; i++) {
        buffer[i] = rfftAry[i].real * rfftAry[i].real + rfftAry[i].imag * rfftAry[i].imag; // not need a sqrtf, the buffer need powf
    }
    buffer[0] = fabs(rfftAry[0].real) * (fabs(rfftAry[0].real));
    buffer[(FRAME_LEN >> 1)] = fabs(rfftAry[0].imag) * (fabs(rfftAry[0].imag));

//    float melSGRet[NUM_FBANK_BINS];
//    memset(melSGRet, 0, NUM_FBANK_BINS * sizeof(float));
    memset(melSG_out, 0, NUM_FBANK_BINS * sizeof(uint8_t));

    float fft_bin_width = 1.0f * SAMP_FREQ / FRAME_LEN;
    float mel_low_freq = MelScale(MEL_LOW_FREQ);
    float mel_high_freq = MelScale(MEL_HIGH_FREQ);
    float mel_freq_delta = (mel_high_freq - mel_low_freq) / (NUM_FBANK_BINS - 1);

    for (int melNdx = 0; melNdx < NUM_FBANK_BINS; melNdx++) {
        float mel = melNdx * mel_freq_delta + mel_low_freq;
        float hz = InverseMelScale(mel);
        float hzNdxFp = hz / fft_bin_width;
        int hzNdx0 = (int)hzNdxFp;
        int hzNdx1 = (int)(ceilf(hzNdxFp));
        if (hzNdx1 == hzNdx0) {
            hzNdx0 -= 1;
        }
        float weit0 = hzNdx1 - hzNdxFp;
        float weit1 = 1 - weit0;
        float melEnergy = buffer[hzNdx0] * weit0 + buffer[hzNdx1] * weit1 + EPSL;
        float mel_dB = log10f(melEnergy) * 20;
        melSG_out[NUM_FBANK_BINS - 1 - melNdx] = quantize_melSG(mel_dB);
    }
    //memcpy(melSG_out, melSGRet, NUM_FBANK_BINS * sizeof(float));
}